"""Federated server implementation."""
import pickle
import torch
from typing import Dict, Any, List, Optional
from .protocols.auth import ECDSAAuth, AuthScheme
from . import model as model_mod
from typing import Optional

class FedServer:
    def __init__(self, auth_scheme: AuthScheme = None, device: str = 'cpu'):
        self.auth = auth_scheme or ECDSAAuth()
        self.clients: Dict[str, Dict[str, Any]] = {}
        self.global_model = model_mod.get_model()
        self.device = device
        self.update_logs: List[Dict[str, Any]] = []  # per update metrics (round filled by experiment layer)

    def register_client(self, client_id: str, pk: Any):
        self.clients[client_id] = { 'pk': pk, 'samples': 0 }

    def receive_update(self, client_id: str, delta_state: Dict[str, torch.Tensor], signature: bytes, metadata: Dict[str, Any], raw_message: Optional[bytes] = None) -> bool:
        """Receive a client update and verify its signature.

        raw_message: original serialized bytes that were signed on client side.
        If not provided, the server re-serializes delta_state (must match client serialization exactly).
        Returns True if accepted, False otherwise.
        """
        message = raw_message if raw_message is not None else pickle.dumps({k: v.cpu() for k,v in delta_state.items()})
        pk = self.clients[client_id]['pk']
        valid, verify_time = self.auth.verify(pk, message, signature)
        if not valid:
            self.update_logs.append({
                'client_id': client_id,
                'accepted': False,
                'samples': metadata.get('samples',0),
                'verify_time': verify_time,
                'signature_size': len(signature),
                'message_size': len(message)
            })
            return False
        samples = metadata.get('samples', 1)
        self.clients[client_id]['pending'] = (delta_state, samples)
        self.clients[client_id]['samples'] = samples
        self.update_logs.append({
            'client_id': client_id,
            'accepted': True,
            'samples': samples,
            'verify_time': verify_time,
            'signature_size': len(signature),
            'message_size': len(message)
        })
        return True

    def aggregate(self):
        # Simple FedAvg on collected pending deltas
        total_samples = sum(c.get('samples',0) for c in self.clients.values() if 'pending' in c)
        global_state = self.global_model.state_dict()
        for cid, entry in self.clients.items():
            if 'pending' not in entry:
                continue
            delta, n = entry['pending']
            weight = n / total_samples if total_samples > 0 else 0
            for k in global_state:
                global_state[k] = global_state[k] + delta[k] * weight
            del entry['pending']
        self.global_model.load_state_dict(global_state)

    def get_global_model_state(self):
        return self.global_model.state_dict()

    def collect_round_metrics(self, round_id: int):
        for row in self.update_logs:
            if 'round' not in row:
                row['round'] = round_id
        rows = self.update_logs
        self.update_logs = []
        return rows
